import torch
from timm import create_model

# 重新创建模型结构（必须与训练时完全一致）
model = create_model("efficientnet_b0", pretrained=False, num_classes=17)  # num_classes 替换为你的类别数
model.load_state_dict(torch.load("efficientnet_b0_custom.pth"))  # 替换为你的模型路径
model.eval()  # 设置为评估模式（关闭Dropout/BatchNorm等）
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
from torchvision import transforms

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
from PIL import Image
import numpy as np


def predict_image(image_path, model, transform):
    img = Image.open(image_path).convert("RGB")
    img_tensor = transform(img).unsqueeze(0)  # 增加batch维度
    img_tensor = img_tensor.to("cuda" if torch.cuda.is_available() else "cpu")

    with torch.no_grad():
        output = model(img_tensor)
        prob = torch.nn.functional.softmax(output[0], dim=0)
        pred_class = torch.argmax(prob).item()

    return pred_class, prob[pred_class].item()


# 使用示例
class_names = ['丝瓜', '人参果', '佛手瓜', '冬瓜', '南瓜', '哈密瓜', '木瓜', '甜瓜-伊丽莎白', '甜瓜-白', '甜瓜-绿', '甜瓜-金', '白兰瓜', '羊角蜜', '苦瓜', '西瓜', '西葫芦', '黄瓜']  # 替换为你的类别名称
# image_path = "/home/featurize/work/melon17_full/test/白兰瓜/105.jpg"
# pred_class, confidence = predict_image(image_path, model, test_transform)
# print(f"Predicted: {class_names[pred_class]} (Confidence: {confidence:.2%})")

from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# 加载测试数据集
test_dataset = ImageFolder("/home/featurize/work/melon17_full/test", transform=test_transform)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# 评估函数
def evaluate(model, loader):
    correct = 0
    total = 0
    model.eval()

    with torch.no_grad():
        for images, labels in loader:
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f"Test Accuracy: {100 * correct / total:.2f}%")


# 运行评估
evaluate(model, test_loader)